package com.cl.ks.flow.handler;

import com.cl.ks.flow.util.StaticCollection;
import com.cl.spark.dto.SparkResult;
import com.cl.ks.entity.KsFlowResult;
import com.cl.ks.flow.base.BaseFlowNodeHandler;
import com.cl.ks.flow.enums.NodeHandlerEnum;
import com.cl.ks.service.KsFlowGraphService;
import com.cl.ks.service.KsFlowResultService;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;

@Component
public class EndHandler extends BaseFlowNodeHandler {
    private final KsFlowGraphService ksFlowGraphService;

    private final KsFlowResultService ksFlowResultService;

    public EndHandler(KsFlowGraphService ksFlowGraphService, KsFlowResultService ksFlowResultService) {
        this.ksFlowGraphService = ksFlowGraphService;
        this.ksFlowResultService = ksFlowResultService;
    }

    @Override
    public SparkResult process(BaseFlowNodeHandler.ProcessParam processParam) {
        String sparkNodeName = "结束";
        Integer ksFlowTaskId = processParam.getKsFlowGraph().getKsFlowTaskId();
        Integer ksFlowGraphRecordId = processParam.getKsFlowGraphRecord().getId();
        String createdBy = processParam.getKsFlowGraph().getCreatedBy();
        String createdByShowName = processParam.getKsFlowGraph().getCreatedByShowName();
        String ruleIds = String.join(",", processParam.getKsFlowGraph().getRuleIdList());
        String sourceConfigIds = String.join(",", processParam.getKsFlowGraph().getSourceConfigIdList());
        String dataTypeCodes = String.join(",", processParam.getKsFlowGraph().getKsDataTypeCodeList());
        String createTables = String.join(",", processParam.getKsFlowGraph().getCreateTableList());

        processParam.getKsFlowGraph().setRuleIds(ruleIds);
        processParam.getKsFlowGraph().setSourceConfigIds(sourceConfigIds);
        processParam.getKsFlowGraph().setKsDataTypeCodes(dataTypeCodes);
        processParam.getKsFlowGraph().setCreateTables(createTables);

        Integer ksflowGraphId = processParam.getKsFlowGraph().getId();
        if (processParam.getKsFlowGraph().getIncrementFlag() == null || !processParam.getKsFlowGraph().getIncrementFlag()) {
            KsFlowResult deleteParam = new KsFlowResult();
            deleteParam.setKsFlowGraphId(ksflowGraphId);
            ksFlowResultService.deleteAll(deleteParam);
        }

        Map<String, SparkResult> sparkResultList = processParam.getSparkNode().getInputSparkResultList();

        StaticCollection.newKsFlowResultList(ksflowGraphId);
        List<KsFlowResult> ksFlowResultList = new ArrayList<>();
        for (Map.Entry<String, SparkResult> entry : sparkResultList.entrySet()) {
            String nodeName = entry.getKey();
            Dataset<Row> dataset = entry.getValue().getDataset();

            dataset.limit(10000).foreach(item -> {
                KsFlowResult ksFlowResult = new KsFlowResult();
                ksFlowResult.setKsFlowGraphRecordId(ksFlowGraphRecordId);
                ksFlowResult.setKsFlowGraphId(ksflowGraphId);
                ksFlowResult.setNodeName(nodeName);
                ksFlowResult.setResult(item.json());
                ksFlowResult.setKsFlowTaskId(ksFlowTaskId);
                ksFlowResult.setCreatedBy(createdBy);
                ksFlowResult.setCreatedByShowName(createdByShowName);
                StaticCollection.getKsFlowResultList(ksflowGraphId).add(ksFlowResult);
            });
        }
        Long count = sparkResultList.values().stream().map(SparkResult::getCount).reduce(Long::sum).orElse(0L);
        try {
            ksFlowResultService.saveAllPartition(StaticCollection.getKsFlowResultList(ksflowGraphId), 1000);
        }catch (Exception e) {
            e.printStackTrace();
        }
        StaticCollection.removeKsFlowResult(ksflowGraphId);
        return SparkResult.success(Objects.requireNonNull(sparkResultList.values().stream().findFirst().orElse(null)).getDataset(), count);
    }

    @Override
    public NodeHandlerEnum getType() {
        return NodeHandlerEnum.END;
    }
}
